import logging
from typing import List, Dict, Any, Optional
import re
import time
import threading
from functools import wraps

logger = logging.getLogger(__name__)

# 超时装饰器
def timeout(seconds=10, default=None):
    """函数执行超时装饰器"""
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            result = [default]
            exception = [None]
            
            def target():
                try:
                    result[0] = func(*args, **kwargs)
                except Exception as e:
                    exception[0] = e
            
            thread = threading.Thread(target=target)
            thread.daemon = True
            thread.start()
            thread.join(seconds)
            
            if thread.is_alive():
                logger.warning(f"函数 {func.__name__} 执行超时（{seconds}秒）")
                return default
            
            if exception[0]:
                raise exception[0]
            
            return result[0]
        return wrapper
    return decorator

class LightweightHyDEGenerator:
    """
    轻量级HyDE生成器，减少对LLM的依赖，提高性能
    """
    
    def __init__(self, llm=None):
        """
        初始化轻量级HyDE生成器
        
        Args:
            llm: 语言模型实例（可选）
        """
        self.llm = llm
        # 预定义的模板，用于生成检索增强词
        self.query_templates = {
            'company_background': [
                "{company_name} 的 {aspect} 相关信息包括",
                "关于 {company_name} 的 {aspect}，招股书显示",
                "{company_name} 在 {aspect} 方面的情况是"
            ],
            'financial_data': [
                "{company_name} {year}年的 {metric} 数据显示",
                "根据财务报表，{company_name} {year}年的 {metric} 为",
                "{company_name} 的 {metric} 相关财务指标包括"
            ],
            'general': [
                "关于 {topic} 的详细信息包括",
                "以下是关于 {topic} 的重要信息",
                "{topic} 的相关内容主要有"
            ]
        }
        
    @timeout(seconds=8, default=[])
    def generate_hypothetical_answers(self, query: str, num_answers: int = 1) -> List[str]:
        """
        为给定查询生成假设答案（超轻量级版本，优先使用模板，必要时才调用LLM）
        
        Args:
            query: 用户查询
            num_answers: 要生成的假设答案数量
            
        Returns:
            假设答案列表
        """
        if not query or num_answers <= 0:
            return []
        
        start_time = time.time()
        logger.info(f"为查询生成 {num_answers} 个检索增强表达式: {query[:100]}...")
        
        hypothetical_answers = []
        
        # 1. 首先尝试使用模板生成增强表达式（无需LLM调用）
        template_answers = self._generate_from_templates(query, max(num_answers, 1))
        if template_answers:
            hypothetical_answers.extend(template_answers[:num_answers])
            logger.info(f"从模板生成了 {len(hypothetical_answers)} 个增强表达式")
        
        # 2. 如果模板生成的答案不足且有LLM可用，使用LLM补充
        remaining = num_answers - len(hypothetical_answers)
        if remaining > 0 and self.llm:
            try:
                llm_answers = self._generate_with_llm(query, remaining)
                hypothetical_answers.extend(llm_answers)
                logger.info(f"从LLM生成了 {len(llm_answers)} 个增强表达式")
            except Exception as e:
                logger.warning(f"使用LLM生成增强表达式失败: {e}")
        
        # 3. 如果仍然没有生成足够的答案，使用简单的查询变换
        remaining = num_answers - len(hypothetical_answers)
        if remaining > 0:
            fallback_answers = self._generate_fallback_expressions(query, remaining)
            hypothetical_answers.extend(fallback_answers)
            logger.info(f"生成了 {len(fallback_answers)} 个备用增强表达式")
        
        # 去重并限制数量
        hypothetical_answers = list(dict.fromkeys(hypothetical_answers))[:num_answers]
        
        logger.info(f"成功生成 {len(hypothetical_answers)} 个检索增强表达式，耗时: {time.time() - start_time:.2f}秒")
        return hypothetical_answers
    
    def _generate_from_templates(self, query: str, num_answers: int) -> List[str]:
        """从预定义模板生成增强表达式"""
        results = []
        
        # 提取查询中的关键信息
        extracted_info = self._extract_query_info(query)
        query_type = extracted_info.get('type', 'general')
        
        # 根据查询类型选择模板
        if query_type in self.query_templates:
            templates = self.query_templates[query_type]
            
            # 生成多个变体
            for i in range(num_answers):
                template = templates[i % len(templates)]
                
                try:
                    # 填充模板
                    if query_type == 'company_background':
                        result = template.format(
                            company_name=extracted_info.get('company_name', ''),
                            aspect=extracted_info.get('aspect', '')
                        )
                    elif query_type == 'financial_data':
                        result = template.format(
                            company_name=extracted_info.get('company_name', ''),
                            year=extracted_info.get('year', ''),
                            metric=extracted_info.get('metric', '')
                        )
                    else:
                        result = template.format(
                            topic=extracted_info.get('topic', query)
                        )
                    
                    # 确保生成的表达式不为空
                    if result and not self._is_invalid_answer(result):
                        results.append(result)
                except KeyError:
                    # 如果缺少必要的信息，跳过此模板
                    continue
        
        return results
    
    def _generate_with_llm(self, query: str, num_answers: int) -> List[str]:
        """使用LLM生成增强表达式（仅在必要时使用）"""
        results = []
        
        # 构建更简洁的提示词，专注于生成检索增强表达式
        prompt = f"""你需要为以下查询生成检索增强表达式，使其更适合向量检索。请保持简洁，突出关键概念。

原始查询: {query}

增强表达式（只需输出表达式本身，不要添加其他说明）:"""
        
        for i in range(num_answers):
            try:
                # 使用更低的温度和更短的token限制，加快生成速度
                response = self.llm.invoke(
                    prompt,
                    temperature=0.5,
                    max_tokens=100,
                    timeout=5  # 设置LLM调用超时
                )
                
                # 提取生成的内容
                if hasattr(response, 'content'):
                    answer = response.content.strip()
                else:
                    answer = str(response).strip()
                
                if answer and not self._is_invalid_answer(answer):
                    results.append(answer)
            except Exception as e:
                logger.error(f"LLM生成第 {i+1} 个增强表达式时出错: {e}")
                break  # 出错时停止尝试
        
        return results
    
    def _generate_fallback_expressions(self, query: str, num_answers: int) -> List[str]:
        """生成备用增强表达式（无需外部依赖）"""
        results = []
        
        # 简单的查询变换策略
        transformations = [
            lambda q: f"{q} 详细信息",
            lambda q: f"关于 {q} 的重要内容",
            lambda q: f"{q} 主要情况",
            lambda q: f"{q} 关键数据"
        ]
        
        for i in range(num_answers):
            transform = transformations[i % len(transformations)]
            result = transform(query)
            results.append(result)
        
        return results
    
    def _extract_query_info(self, query: str) -> Dict[str, str]:
        """从查询中提取关键信息"""
        info = {'type': 'general'}
        
        # 提取年份信息
        year_match = re.search(r'20\d{2}|19\d{2}', query)
        if year_match:
            info['year'] = year_match.group()
        
        # 提取常见财务指标
        financial_metrics = ['毛利率', '收入', '占比', '净利润', '收益率', '市盈率', '市净率']
        for metric in financial_metrics:
            if metric in query:
                info['metric'] = metric
                info['type'] = 'financial_data'
                break
        
        # 检查是否包含公司相关关键词
        company_aspects = ['股权结构', '代持', '股东', '持股', '控股股东', '历史沿革']
        for aspect in company_aspects:
            if aspect in query:
                info['aspect'] = aspect
                info['type'] = 'company_background'
                # 尝试提取公司名称（简单方法）
                company_name = self._extract_company_name(query)
                if company_name:
                    info['company_name'] = company_name
                break
        
        # 如果没有识别出特定类型，使用通用主题
        if info['type'] == 'general':
            info['topic'] = query
        
        return info
    
    def _extract_company_name(self, query: str) -> Optional[str]:
        """简单提取公司名称"""
        # 这是一个简化版本，实际应用中可能需要更复杂的命名实体识别
        company_patterns = [
            r'[\u4e00-\u9fa5]+(股份有限公司|有限公司|集团|公司)',
            r'[\u4e00-\u9fa5]+科技',
            r'[\u4e00-\u9fa5]+集团'
        ]
        
        for pattern in company_patterns:
            match = re.search(pattern, query)
            if match:
                return match.group()
        
        return None
    
    def _is_invalid_answer(self, answer: str) -> bool:
        """
        检查答案是否无效
        """
        # 检查是否为空
        if not answer:
            return True
        
        # 检查是否过于简短
        if len(answer) < 10:
            return True
        
        return False

class LightweightHyDEEnhancer:
    """
    轻量级HyDE增强器，减少资源消耗，提高性能
    """
    
    def __init__(self, hyde_generator: LightweightHyDEGenerator):
        """
        初始化轻量级HyDE增强器
        
        Args:
            hyde_generator: 轻量级HyDE生成器实例
        """
        self.hyde_generator = hyde_generator
    
    @timeout(seconds=15, default=[])
    def enhance_query(self, original_query: str, vector_store, company_name: str = None, k: int = 4) -> List[Any]:
        """
        使用轻量级HyDE增强查询，减少资源消耗，提高性能
        
        Args:
            original_query: 原始用户查询
            vector_store: 向量存储实例
            company_name: 公司名称（可选）
            k: 要返回的结果数量
            
        Returns:
            增强检索的文档列表
        """
        start_time = time.time()
        logger.info(f"使用轻量级HyDE增强查询: {original_query[:100]}...")
        
        # 生成检索增强表达式（替代完整的假设答案）
        enhanced_expressions = self.hyde_generator.generate_hypothetical_answers(original_query, num_answers=2)
        
        # 如果没有生成有效的增强表达式，回退到原始查询
        if not enhanced_expressions:
            logger.warning("未能生成有效的检索增强表达式，使用原始查询进行检索")
            if company_name:
                return vector_store.company_specific_search(company_name, original_query, k=k)
            else:
                return vector_store.similarity_search(original_query, k=k)
        
        # 策略1: 首先尝试使用简单的关键词增强（消耗最小）
        try:
            # 提取关键词
            query_keywords = self._extract_keywords(original_query)
            if query_keywords:
                # 构建基于关键词的增强查询
                keyword_enhanced_query = original_query + " " + " ".join(query_keywords)
                
                logger.info(f"使用关键词增强查询")
                if company_name:
                    docs = vector_store.company_specific_search(company_name, keyword_enhanced_query, k=k)
                else:
                    docs = vector_store.similarity_search(keyword_enhanced_query, k=k)
                
                # 如果检索到结果，直接返回
                if docs:
                    logger.info(f"关键词增强检索成功，返回 {len(docs)} 个结果，耗时: {time.time() - start_time:.2f}秒")
                    return docs
        except Exception as e:
            logger.warning(f"关键词增强检索出错: {e}")
        
        # 策略2: 使用生成的增强表达式进行检索
        try:
            # 合并原始查询和增强表达式
            combined_query = original_query
            for expr in enhanced_expressions:
                combined_query += " " + expr
                
                # 限制总长度，防止查询过长
                if len(combined_query) > 300:  # 进一步缩短查询长度
                    combined_query = combined_query[:300]
                    break
            
            logger.info(f"使用增强表达式进行检索")
            if company_name:
                docs = vector_store.company_specific_search(company_name, combined_query, k=min(k*2, 6))  # 减少返回数量以提高速度
            else:
                docs = vector_store.similarity_search(combined_query, k=min(k*2, 6))
            
            # 如果检索到结果，进行简单排序后返回
            if docs:
                # 使用简化的排序逻辑，减少计算量
                sorted_docs = self._simple_sort_docs(docs, original_query, company_name)
                final_docs = sorted_docs[:k]
                logger.info(f"增强表达式检索成功，返回 {len(final_docs)} 个结果，耗时: {time.time() - start_time:.2f}秒")
                return final_docs
        except Exception as e:
            logger.warning(f"增强表达式检索出错: {e}")
        
        # 最终回退到原始查询
        logger.warning("所有增强策略失败，回退到原始查询")
        if company_name:
            docs = vector_store.company_specific_search(company_name, original_query, k=k)
        else:
            docs = vector_store.similarity_search(original_query, k=k)
        
        logger.info(f"HyDE增强检索完成，返回 {len(docs)} 个结果，耗时: {time.time() - start_time:.2f}秒")
        return docs
    
    def _simple_sort_docs(self, docs: List[Any], original_query: str, company_name: str = None) -> List[Any]:
        """
        简化的文档排序函数，减少计算量
        """
        if not docs:
            return []
        
        # 提取原始查询的核心关键词
        query_keywords = self._extract_keywords(original_query, max_keywords=3)  # 只提取3个核心关键词
        
        scored_docs = []
        for doc in docs:
            score = 0
            content = getattr(doc, 'page_content', str(doc)).lower()
            
            # 公司名称匹配加分（如果有）
            if company_name and company_name.lower() in content:
                score += 1.0
            
            # 关键词匹配加分（简化版，只检查是否存在）
            for keyword in query_keywords:
                if keyword.lower() in content:
                    score += 0.5
                    break  # 每个文档只加一次关键词分，减少计算量
            
            # 简单的长度过滤（不加分，只过滤）
            content_length = len(content)
            if content_length < 100:
                score -= 0.5  # 太短的文档可能质量不高
            
            scored_docs.append((doc, score))
        
        # 按分数排序，降序
        scored_docs.sort(key=lambda x: x[1], reverse=True)
        
        # 返回排序后的文档列表（去掉分数）
        return [doc for doc, score in scored_docs]
    
    def _extract_keywords(self, text: str, max_keywords: int = 5) -> List[str]:
        """
        简化的关键词提取函数，减少计算量
        """
        # 移除常见的停用词
        stop_words = ['的', '了', '在', '是', '我', '有', '和', '就', '不', '人', '都', '一', '一个', '上', '也', '很', '到', '说', '要', '去', '你', '会', '着', '没有', '看', '好', '自己', '这']
        
        # 提取关键词（过滤掉停用词和标点符号）
        keywords = []
        words = re.findall(r'[\u4e00-\u9fa5a-zA-Z]+', text)  # 提取中文字符和英文字母
        
        for word in words:
            if word not in stop_words and len(word) > 1:
                keywords.append(word)
                if len(keywords) >= max_keywords:
                    break  # 达到最大关键词数量后停止
        
        return keywords

# 工厂函数，用于创建轻量级HyDE组件
def create_hyde_enhancer():
    """
    创建轻量级HyDE增强器实例（高度优化版）
    """
    start_time = time.time()
    try:
        # 尝试获取LLM实例（可选）
        llm = None
        try:
            from app.agent.llm import get_llm
            llm = get_llm()
            logger.info("成功获取LLM实例")
        except Exception as llm_err:
            logger.warning(f"无法获取LLM实例，将使用模板和关键词增强: {llm_err}")
        
        # 创建轻量级HyDE生成器
        hyde_generator = LightweightHyDEGenerator(llm)
        # 创建轻量级HyDE增强器
        hyde_enhancer = LightweightHyDEEnhancer(hyde_generator)
        
        logger.info(f"轻量级HyDE增强器创建成功，耗时: {time.time() - start_time:.2f}秒")
        return hyde_enhancer
    except Exception as e:
        logger.error(f"创建HyDE增强器失败: {e}")
        import traceback
        logger.debug(f"创建HyDE增强器失败的详细堆栈: {traceback.format_exc()}")
        return None

# 兼容性导入，确保现有代码可以正常工作
HyDEGenerator = LightweightHyDEGenerator
HyDEEnhancer = LightweightHyDEEnhancer